#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>

// 声明任务切换函数（在汇编文件中实现）
extern void task_switch(void* next);

// 定义栈帧结构，与汇编代码中的保存/恢复顺序对应
typedef struct {
    uint32_t edi;        // 保存的寄存器
    uint32_t esi;        // 保存的寄存器
    uint32_t ebx;        // 保存的寄存器
    uint32_t ebp;        // 保存的帧指针
    uint32_t ret_addr;   // 返回地址
} StackFrame;

// 任务控制块结构 - 确保总大小正好是4KB
typedef struct {
    void* esp;                       // 栈指针
    uint8_t stack[4096 - sizeof(void*)]; // 调整栈大小使总大小正好是4KB
} __attribute__((aligned(4096))) TCB;

// 两个任务的TCB
TCB task1;
TCB task2;

// 任务1的函数
void task1_func() {
    int counter = 0;
    while (1) {
        printf("Task 1: %d\n", counter++);
        if (counter % 3 == 0) {
            printf("Task 1: Switching to Task 2\n");
            task_switch(&task2);
        }
    }
}

// 任务2的函数
void task2_func() {
    int counter = 0;
    while (1) {
        printf("Task 2: %d\n", counter++);
        if (counter % 5 == 0) {
            printf("Task 2: Switching to Task 1\n");
            task_switch(&task1);
        }
    }
}

// 初始化任务栈 - 使用结构体表示栈帧
void init_task(TCB* tcb, void (*func)()) {
    // 计算栈顶位置（栈从高地址向低地址增长）
    uintptr_t stack_top = (uintptr_t)(tcb->stack + sizeof(tcb->stack));
    
    // 确保栈顶是16字节对齐（更好的对齐可以提高性能）
    stack_top &= ~0xF;  // 16字节对齐
    
    // 为栈帧分配空间
    stack_top -= sizeof(StackFrame);
    
    // 初始化栈帧
    StackFrame* frame = (StackFrame*)stack_top;
    frame->edi = 0;
    frame->esi = 0;
    frame->ebx = 0;
    frame->ebp = 0;
    frame->ret_addr = (uint32_t)func;
    
    // 设置初始栈指针
    tcb->esp = frame;
    
    // 验证对齐
    if ((uintptr_t)tcb->esp % 4 != 0) {
        printf("警告: 栈指针未4字节对齐! ESP = %p\n", tcb->esp);
    } else {
        printf("栈指针已正确对齐: ESP = %p\n", tcb->esp);
    }
}

// 打印任务信息的辅助函数
void print_task_info(TCB* tcb, const char* name) {
    printf("任务 %s:\n", name);
    printf("  TCB 地址: %p\n", (void*)tcb);
    printf("  ESP 指针: %p\n", tcb->esp);
    
    StackFrame* frame = (StackFrame*)tcb->esp;
    printf("  栈帧内容:\n");
    printf("    EDI: 0x%08X\n", frame->edi);
    printf("    ESI: 0x%08X\n", frame->esi);
    printf("    EBX: 0x%08X\n", frame->ebx);
    printf("    EBP: 0x%08X\n", frame->ebp);
    printf("    返回地址: %p\n", (void*)frame->ret_addr);
    
    // 验证栈对齐
    printf("  栈对齐: %s (余数: %lu)\n", 
           (uintptr_t)tcb->esp % 4 == 0 ? "正确" : "错误",
           (uintptr_t)tcb->esp % 4);
}

int main() {
    // 初始化任务
    init_task(&task1, task1_func);
    init_task(&task2, task2_func);
    
    // 打印任务信息
    printf("=== 任务初始化完成 ===\n");
    print_task_info(&task1, "Task 1");
    print_task_info(&task2, "Task 2");
    
    printf("\n按任意键开始执行 Task 1...\n\n");
    
    //an key to start
    getchar();

    // 开始执行任务1
    task_switch(&task1);
    
    // 不应该返回到这里
    printf("错误: 返回到main函数\n");
    return 0;
}
